from composite_reward_cal import CompositeRewardCalculator

# 导入所有奖励组件
from RL.Reword.impl import (
    ComplianceRewardComponent, 
    TimeEfficiencyComponent,
    ResourceUtilizationComponent,
    BusinessGoalComponent,
    CostReductionComponent,
    CustomerSatisfactionComponent
)

# 创建奖励计算器
reward_calculator = CompositeRewardCalculator()

# 注册标准组件
reward_calculator.register_component('compliance', ComplianceRewardComponent(), 0.4)
reward_calculator.register_component('time', TimeEfficiencyComponent(coefficient=-0.1), 0.3)
reward_calculator.register_component('resource', ResourceUtilizationComponent(coefficient=0.5), 0.2)
reward_calculator.register_component('business', BusinessGoalComponent(goal_reward=2.0), 0.1)

# 注册业务特定组件
cost_data = {'申请提交': 10, '材料审核': 20, '初审': 30, '终审': 50, '审批通过': 15, '审批拒绝': 10}
reward_calculator.register_component('cost', CostReductionComponent(cost_data=cost_data), 0.15)

satisfaction_data = {'申请提交': 7, '材料审核': 5, '初审': 6, '终审': 4, '审批通过': 9, '审批拒绝': 3}
reward_calculator.register_component('satisfaction', CustomerSatisfactionComponent(satisfaction_data=satisfaction_data), 0.1)

# 调整权重
reward_calculator.set_weight('compliance', 0.3)
reward_calculator.set_weight('cost', 0.25)

# 使用奖励计算器替代原始函数
def composite_reward(current_activity, next_activity, case_events, context=None):
    """使用组件化的复合奖励计算器"""
    context = context or {}  # 创建上下文对象用于传递额外信息
    reward = reward_calculator.calculate_reward(
        current_activity, next_activity, case_events, context
    )
    
    # 可以访问并记录详细的奖励分解
    if 'reward_breakdown' in context:
        print("奖励明细:", context['reward_breakdown'])
        
    return reward